/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "bsr_common.h"
#include "parbsr_smooth.h"
#ifdef USE_OPENMP
#include <omp.h>
#endif

namespace YHAMG
{

void ParBSR_Smooth(const ParBSR_Matrix& A, const double* D_LU, const Par_Vector& b, const Par_Vector& x, const Par_Vector& w, int type, double omega)
{
	int n = A.LocLoc.size[0];
	int Ext_size = A.LocExt.size[1];

	int block_size = A.LocLoc.block_size;
	int block_nnz = block_size * block_size;

	int* LocLoc_rowptr = A.LocLoc.rowptr;
	int* LocLoc_colind = A.LocLoc.colind;
	double* LocLoc_values = A.LocLoc.values;
	int* LocExt_rowptr = A.LocExt.rowptr;
	int* LocExt_colind = A.LocExt.colind;
	double* LocExt_values = A.LocExt.values;

	double* xv_Loc = x.Loc.values;
	double* bv_Loc = b.Loc.values;
	double* wv_Loc = w.Loc.values;

	double* xv_Ext = A.x_Ext.values;

	A.ExchangeHalo(x.Loc, A.x_Ext);

	if (type == 0)
	{
		for (int i = 0; i < n; ++i)
			VecBlockCopy(block_size, xv_Loc + i * block_size, wv_Loc + i * block_size);
#ifdef USE_OPENMP
#pragma omp parallel
#endif
		{
			double* temp = new double[block_size];
#ifdef USE_OPENMP
#pragma omp for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
				for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, wv_Loc + LocLoc_colind[j] * block_size, temp);
				for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
					BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
				BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
				VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
			}

			delete[] temp;
		}
	}
	else if (type == 1)
	{
		int num_threads = 1;
#ifdef USE_OPENMP
		num_threads = omp_get_max_threads();
#endif
		if (num_threads == 1)
		{
			double* temp = new double[block_size];

			for (int i = 0; i < n; ++i)
			{
				VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
				for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, xv_Loc + LocLoc_colind[j] * block_size, temp);
				for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
					BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
				BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
				VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
			}

			delete[] temp;
		}
		else
		{
			for (int i = 0; i < n; ++i)
				VecBlockCopy(block_size, xv_Loc + i * block_size, wv_Loc + i * block_size);
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
			for (int t = 0; t < num_threads; ++t)
			{
				int begin = n * t / num_threads;
				int end = n * (t + 1) / num_threads;
				double* temp = new double[block_size];

				for (int i = begin; i < end; ++i)
				{
					VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
					for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					{
						if (LocLoc_colind[j] < begin || LocLoc_colind[j] >= end)
							BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, wv_Loc + LocLoc_colind[j] * block_size, temp);
						else
							BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, xv_Loc + LocLoc_colind[j] * block_size, temp);
					}
					for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
						BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
					BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
					VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
				}

				delete[] temp;
			}
		}
	}
	else if (type == 2)
	{
		int num_threads = 1;
#ifdef USE_OPENMP
		num_threads = omp_get_max_threads();
#endif

		if (num_threads == 1)
		{
			double* temp = new double[block_size];

			for (int i = n - 1; i >= 0; --i)
			{
				VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
				for (int j = LocLoc_rowptr[i + 1] - 1; j >= LocLoc_rowptr[i]; --j)
					BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, xv_Loc + LocLoc_colind[j] * block_size, temp);
				for (int j = LocExt_rowptr[i + 1] - 1; j >= LocExt_rowptr[i]; --j)
					BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
				BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
				VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
			}

			delete[] temp;
		}
		else
		{
			for (int i = 0; i < n; ++i)
				VecBlockCopy(block_size, xv_Loc + i * block_size, wv_Loc + i * block_size);
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
			for (int t = 0; t < num_threads; ++t)
			{
				int begin = n * t / num_threads;
				int end = n * (t + 1) / num_threads;
				double* temp = new double[block_size];

				for (int i = end - 1; i >= begin; --i)
				{
					VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
					for (int j = LocLoc_rowptr[i + 1] - 1; j >= LocLoc_rowptr[i]; --j)
					{
						if (LocLoc_colind[j] < begin || LocLoc_colind[j] >= end)
							BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, wv_Loc + LocLoc_colind[j] * block_size, temp);
						else
							BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, xv_Loc + LocLoc_colind[j] * block_size, temp);
					}
					for (int j = LocExt_rowptr[i + 1] - 1; j >= LocExt_rowptr[i]; --j)
						BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
					BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
					VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
				}

				delete[] temp;
			}
		}
	}
	else if (type == 3)
	{
		int num_threads = 1;
#ifdef USE_OPENMP
		num_threads = omp_get_max_threads();
#endif

		if (num_threads == 1)
		{
			double* temp = new double[block_size];

			for (int i = 0; i < n; ++i)
			{
				VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
				for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, xv_Loc + LocLoc_colind[j] * block_size, temp);
				for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
					BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
				BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
				VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
			}

			for (int i = n - 1; i >= 0; --i)
			{
				VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
				for (int j = LocLoc_rowptr[i + 1] - 1; j >= LocLoc_rowptr[i]; --j)
					BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, xv_Loc + LocLoc_colind[j] * block_size, temp);
				for (int j = LocExt_rowptr[i + 1] - 1; j >= LocExt_rowptr[i]; --j)
					BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
				BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
				VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
			}

			delete[] temp;
		}
		else
		{
			for (int i = 0; i < n; ++i)
				VecBlockCopy(block_size, xv_Loc + i * block_size, wv_Loc + i * block_size);
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
			for (int t = 0; t < num_threads; ++t)
			{
				int begin = n * t / num_threads;
				int end = n * (t + 1) / num_threads;
				double* temp = new double[block_size];

				for (int i = begin; i < end; ++i)
				{
					VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
					for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					{
						if (LocLoc_colind[j] < begin || LocLoc_colind[j] >= end)
							BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, wv_Loc + LocLoc_colind[j] * block_size, temp);
						else
							BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, xv_Loc + LocLoc_colind[j] * block_size, temp);
					}
					for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
						BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
					BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
					VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
				}

				for (int i = end - 1; i >= begin; --i)
				{
					VecBlockCopy(block_size, bv_Loc + i * block_size, temp);
					for (int j = LocLoc_rowptr[i + 1] - 1; j >= LocLoc_rowptr[i]; --j)
					{
						if (LocLoc_colind[j] < begin || LocLoc_colind[j] >= end)
							BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, wv_Loc + LocLoc_colind[j] * block_size, temp);
						else
							BSRBlockMatVecSub(block_size, LocLoc_values + j * block_nnz, xv_Loc + LocLoc_colind[j] * block_size, temp);
					}
					for (int j = LocExt_rowptr[i + 1] - 1; j >= LocExt_rowptr[i]; --j)
						BSRBlockMatVecSub(block_size, LocExt_values + j * block_nnz, xv_Ext + LocExt_colind[j] * block_size, temp);
					BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, temp);
					VecBlockScaleAdd(block_size, omega, temp, xv_Loc + i * block_size);
				}

				delete[] temp;
			}
		}
	}
}

}